use std::cmp::max;

use primitive_types::U256;

pub fn find_max_prime_factor(number: u128) -> u128 {
    println!("INPUT: {number}");
    let mut max_factor = 0;
    calculate_max_factor(number, &mut max_factor);
    println!("ANS: {max_factor}");
    max_factor
}

/// 取绝对值.
fn big_abs(a: u128, b: u128) -> u128 {
    if a >= b {
        a - b
    } else {
        b - a
    }
}

/// 带取模的乘. a * b % modulus.
fn mul_mod(a: u128, b: u128, modulus: u128) -> u128 {
    let (a, b, modulus) = (U256::from(a), U256::from(b), U256::from(modulus));
    (a * b % modulus).as_u128()
}

/// 带取模的快速幂. base^exponent % modulus.
///
/// 把exponent以二进制形式拆开. 结果为$\prod base^{2^exponent_i}$.
fn quick_pow_mod(mut base: u128, mut exponent: u128, modulus: u128) -> u128 {
    // let (base, exponent, modulus) = (U256::from(base), U256::from(exponent), U256::from(modulus));
    let mut ans = 1;
    while exponent > 0 {
        if exponent % 2 > 0 {
            ans = mul_mod(ans, base, modulus);
        }
        base = mul_mod(base, base, modulus);
        exponent /= 2;
    }
    ans
}

/// 求最大公约数.
fn gcd(a: u128, b: u128) -> u128 {
    if b == 0 {
        a
    } else {
        gcd(b, a % b)
    }
}

/// Miller–Rabin 素性测试.
fn miller_rabin(n: u128) -> bool {
    if n < 3 || n % 2 == 0 {
        return n == 2;
    }
    if n % 3 == 0 {
        return n == 3;
    }

    let mut u = n - 1;
    let mut t = 0;
    while u % 2 == 0 {
        u /= 2;
        t += 1;
    }

    const MR_TEST_TIME: u32 = 10;
    for _ in 0..MR_TEST_TIME {
        // a in [2, n-2]
        let a = fastrand::u128(2..=(n - 2));
        let mut v = quick_pow_mod(a, u, n);
        if v == 1 {
            continue;
        }

        let mut pass = false;
        for _ in 0..t {
            if v == n - 1 {
                pass = true;
                break;
            }
            v = mul_mod(v, v, n);
        }
        if !pass {
            return false;
        }
    }
    return true;
}

/// Pollard-Rho 寻找非平凡因子.
fn pollard_rho(n: u128) -> Option<u128> {
    // f(x) = x^2 + c mod n
    let c = fastrand::u128(1..n);
    let fx = |x: u128| {
        let x_sq = mul_mod(x, x, n);
        let value = (U256::from(x_sq) + c) % n;
        value.as_u128()
    };

    // Brent 判环
    let mut slow: u128 = 0;
    let mut quick: u128 = 0;
    // 每轮快指针走的次数. 指数增长
    let mut ptr_goal = 1;
    loop {
        // 累积乘积
        let mut gathered_val: u128 = 1;
        for step in 1..=ptr_goal {
            quick = fx(quick);
            gathered_val = mul_mod(gathered_val, big_abs(quick, slow), n);
            // 指针相遇
            if gathered_val == 0 {
                return None;
            }

            // 倍增优化
            const GCD_ROUND: u32 = 127;
            if step % GCD_ROUND == 0 {
                let factor = gcd(gathered_val, n);
                if factor > 1 {
                    // 非平凡因子
                    return Some(factor);
                }
            }
        }

        // 倍增优化剩余乘积
        let factor = gcd(gathered_val, n);
        if factor > 1 {
            // 非平凡因子
            return Some(factor);
        }

        ptr_goal *= 2;
        slow = quick;
    }
}

/// 递归求最大素因子.
fn calculate_max_factor(mut n: u128, max_factor: &mut u128) {
    // 不可能有新解, 截断
    if n <= *max_factor || n < 2 {
        return;
    }

    // 素因子
    if miller_rabin(n) {
        *max_factor = max(*max_factor, n);
        return;
    }

    let p;
    loop {
        if let Some(factor) = pollard_rho(n) {
            p = factor;
            break;
        }
    }
    // 消去n中的p因子
    while n % p == 0 {
        n /= p;
    }
    // 递归
    calculate_max_factor(n, max_factor);
    calculate_max_factor(p, max_factor);
}

// mod tests {
//     use super::*;

//     #[test]
//     fn test_miller_rabin() {
//         assert!(miller_rabin(2));
//         assert!(miller_rabin(3));
//         assert!(miller_rabin(5));
//         assert!(!miller_rabin(21));
//         assert!(miller_rabin(41));
//         assert!(!miller_rabin(65535));
//         assert!(miller_rabin(65537));
//         assert!(miller_rabin(6700417));
//         assert!(miller_rabin(67280421310721));
//         assert!(!miller_rabin(340282366920938463463374607431768211455));
//     }

//     #[test]
//     fn test_find_max_prime_factor() {
//         assert_eq!(
//             find_max_prime_factor(340282366920938463463374607431768211455),
//             67280421310721
//         );
//     }
// }
